package testrank.predictors;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import testrank.Conflict;
import testrank.Predictor;

/**
 * @author Hagai Cibulski
 *
 */
public class AffinityPropagationPredictor implements Predictor {

	private static final long serialVersionUID = 1L;
	
	private static final int N_ITERATIONS = 3;
	
	private static final double PROPAGATION_FACTOR = 0.25;

	private AffinityPredictor affinityPredictor;
	
	private CallCountPredictor callCountPredictor;
	
	protected Map<Conflict, Double> conflictAffinities = new HashMap<Conflict, Double>();
	
	// test->method->(affinity, callersList, calleesList)
	protected transient Map<String, Map<String, AffinityVertex>> affinityGraph;
	
	public AffinityPropagationPredictor() {
	}
	
	
	public AffinityPropagationPredictor(AffinityPredictor affinityPredictor,
			CallCountPredictor callCountPredictor) {
		this.affinityPredictor = affinityPredictor;
		this.callCountPredictor = callCountPredictor;
	}


	@Override
	public double rankFailure(String methodSignature, String testSignature) {
		Double affinity = conflictAffinities.get(new Conflict(methodSignature, testSignature));
		if (affinity == null)
			return 0;
		
		assert !Double.isNaN(affinity); 
		return affinity;
	}
	
	public void process() {

		createConflictAffinities();
		
		// create test->method->(affinity, callersList, calleesList)
		affinityGraph = new HashMap<String, Map<String,AffinityVertex>>();
		for (Conflict conflict : conflictAffinities.keySet()) {
			initAffinityGraphVertex(conflict);
		}
		
		// add parents and children from call count predictor 
		for (String test :affinityGraph.keySet()) {
			addEdges(test);
		}
		
		for (int i = 0; i < N_ITERATIONS; i++) {
			propagate();
		}
		
		updateConflictAffinities();
	}


	/**
	 * adds parents and children from call count predictor 
	 */
	private void addEdges(String test) {
		Map<String, AffinityVertex> method2vertex = affinityGraph.get(test);
		for (String method : method2vertex.keySet()) {
			AffinityVertex v = method2vertex.get(method);
			Conflict conflict = new Conflict(method, test);
			if (method.equals(test)) {
				continue;
			}
			Map<String, Integer> calls = callCountPredictor.conflicts.get(conflict);
			Set<String> callers = calls.keySet();
			for (String caller : callers) {
				AffinityVertex parent = null;
				if (caller != null && ! caller.contains("Native Method") && ! caller.contains("Unknown Source")) {
					//Since we get the caller now directly from after-call advice
					String canonicalCaller = caller;
					//String canonicalCaller = StackTraceUtils.toCanonicalSignature(caller);
					if (canonicalCaller != null) {
						parent = method2vertex.get(canonicalCaller);
					}
				}
				if (parent == null) { 
					//TODO assert false;// should never happen after handling ctors in aspect
					parent = new AffinityVertex(0);
					parent.dummy = true;
				}
				Integer weight = calls.get(caller);
				v.addParent(parent, weight);
				parent.addChild(v, weight);
			}
		}
	}


	/**
	 * @param conflict
	 */
	private void initAffinityGraphVertex(Conflict conflict) {
		String test = conflict.getTestSignature();
		String method = conflict.getMethodSignature();
		Double affinity = conflictAffinities.get(conflict);
		
		Map<String, AffinityVertex> method2vertex = affinityGraph.get(test);//conflict);
		if (method2vertex == null) {
			method2vertex = new HashMap<String, AffinityVertex>();
			affinityGraph.put(test, method2vertex);
		}
		AffinityVertex v = new AffinityVertex(affinity);
		
		method2vertex.put(method, v);
	}
	
	private void propagate() {
		// for each test
		for (String test : affinityGraph.keySet()) {
			Map<String, AffinityVertex> method2vertex = affinityGraph.get(test);
			// calculate new generation from previous one without changing previous one 
			for (AffinityVertex v : method2vertex.values()) {
				v.newAffinity =
					PROPAGATION_FACTOR * sumEdges(v) +
					(1- PROPAGATION_FACTOR) * v.affinity;
			}
			// commit new generation
			for (AffinityVertex v : method2vertex.values()) {
				v.affinity = v.newAffinity;
			}
		}
	}

	private double sumEdges(AffinityVertex v) {

		double s= 0.0;
		double max = v.affinity;
		int n = 0;
		
//		// from parents (NOT IN USE)
//		for (AffinityVertex u : v.getParents()) {
//			double a = u.dummy ? v.affinity : u.affinity;
//			double w = 1; 
//			s += a * w;
//			if (a > max) {
//				max = a;
//			}
//			n += w; 
//		}
		// from children
		for (AffinityVertex u : v.getChildren()) {
			double a = u.dummy ? v.affinity : u.affinity;
			double w = 1;
			
			// weights a-la PageRank
//			double repricoralWeight = 1.0 / u.parents.size();
//			double callWeight = v.childrenWeights.get(u);
//			w = callWeight;
//			w = Math.log(callWeight);
//			w = repricoralWeight;
//			w = callWeight * repricoralWeight;
//			w = Math.log(callWeight) * repricoralWeight;
//			w = Math.log(repricoralWeight);
			
			s += a * w;
			if (a > max) {
				max = a;
			}
			n += w; 
		}

		//return max;
		return n == 0 ? 0.0 : s/n;
	}

	private void createConflictAffinities() {
		conflictAffinities = new HashMap<Conflict, Double>();
		for (Conflict conflict : affinityPredictor.conflictAffinities.keySet()) {
			conflictAffinities.put(conflict, affinityPredictor.conflictAffinities.get(conflict));

			// add reflexive affinity for test with itself
			conflictAffinities.put(new Conflict(conflict.getTestSignature(), conflict.getTestSignature()), 1.0);
		}	
	}


	@Override
	public void begin(String methodSignature, String testSignature) {
	}
	@Override
	public void end(String methodSignature, String testSignature) {
	}

	@Override
	public String getName() {
		return "AFFPROP";
	}
	
	private static class AffinityVertex implements Serializable {
		
		private static final long serialVersionUID = 1L;
		
		boolean dummy = false;
		
		double affinity = 0.0;
		double newAffinity = 0.0;
		Map<AffinityVertex, Integer> parentWeights 	= new HashMap<AffinityVertex, Integer>();
		Map<AffinityVertex, Integer> childrenWeights 	= new HashMap<AffinityVertex, Integer>();
		List<AffinityVertex> parents = new ArrayList<AffinityVertex>();
		List<AffinityVertex> children = new ArrayList<AffinityVertex>();

		public AffinityVertex(double affinity) {
			this.affinity = affinity;
		}

		public void addParent(AffinityVertex parent, Integer weight) {
			parentWeights.put(parent, weight);
			parents.add(parent);
		}
			
		public void addChild(AffinityVertex child, Integer weight) {
			childrenWeights.put(child, weight);
			children.add(child);
		}

		public List<AffinityVertex> getChildren() {
			return children;
		}
		
		@SuppressWarnings("unused")
		public List<AffinityVertex> getParents() {
			return parents;
		}
	}
	
	private void updateConflictAffinities() {
		for (Conflict c: conflictAffinities.keySet()) {
			Map<String, AffinityVertex> method2vertex = affinityGraph.get(c.getTestSignature());
			AffinityVertex v = method2vertex.get(c.getMethodSignature());
			assert v != null;	
			conflictAffinities.put(c, v.affinity);
		}
	}
}
